Greptile Summary

This PR fixes a cache-collision bug in get_full_cu_seqlens where two logically distinct tensors — one created under torch.inference_mode() and one in normal training, or one on cuda:0 and another on cuda:1 — could resolve to the same cache key (batch_size, max_seqlen) and incorrectly share a cached cu_seqlens tensor. The fix expands the key to (batch_size, max_seqlen, device, is_inference_mode_enabled()), which is the minimal correct fix for the described Knowledge Distillation scenario.

• Core fix (utils.py): Two lines — capture torch.is_inference_mode_enabled() and build a 4-tuple cache key. No logic beyond that is touched.
• Device isolation: tensor.device is always a concrete torch.device (e.g. cuda:0) so equality and hashing are well-defined; no ambiguity from un-indexed "cuda" strings.
• Inference vs. no-grad: The fix correctly targets torch.inference_mode() (tensors are flagged is_inference=True and cannot leave the context) rather than the weaker torch.no_grad(), which was the actual source of the runtime error in KD training.
• Cache unboundedness: The global _cu_seqlens_cache dict is still never evicted. The new dimensions add at most 2 × num_devices entries per (batch_size, max_seqlen) pair, which is negligible in practice and was already a pre-existing concern.
• Tests: Two new well-structured tests cover both isolation axes; the autouse fixture correctly clears the cache before and after each test to prevent cross-test contamination.

Confidence Score: 5/5

Safe to merge — minimal, targeted fix with no regressions and direct test coverage.

No P0/P1 findings. The change is a two-line key extension that directly addresses the described bug. The torch.device and boolean values used as key components are correctly hashable and produce stable equality. Tests cover both scenarios mentioned in the PR description.

No files require special attention.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["get_full_cu_seqlens called"] --> B{ONNX export mode?}
    B -- Yes --> C["Skip cache, return directly"]
    B -- No --> D["Read torch.is_inference_mode_enabled"]
    D --> E["Build 4-tuple cache lookup: batch+seqlen+device+inference_flag"]
    E --> F{Found in cache?}
    F -- Yes --> G["Return cached cu_seqlens tensor"]
    F -- No --> H["Create new tensor via torch.arange"]
    H --> I["Store in cache and return"]

_{Reviews (7): Last reviewed commit: "Merge branch 'main' into fix/get_full_cu..." | Re-trigger Greptile}

@cyanguwa When you have a moment, could you please take a look at this PR? Thanks:)

@cyanguwa This PR is pretty straightforward. Would you mind taking a quick look? Thank you:)

@cyanguwa Hi:) could you look into this PR? thank you.